from Crypto.PublicKey import RSA
from PIL import Image
import base64
from Crypto.Cipher import PKCS1_v1_5 as PKCS1_cipher


default_encoding = "utf-8"


def mod(x, y):
    return x % y


def bin_ord(flag):
    string = ""
    with open(flag) as f:
        txt = f.read()
        for i in range(len(txt)):
            string = string + bin(ord(txt[i])).replace('0b', '').zfill(8)
    return string


def hide(pic, flag, new_pic):
    count = 0
    im = Image.open(pic)
    width = im.size[0]
    height = im.size[1]
    string = bin_ord(flag)
    for h in range(height):
        for w in range(width):
            pixel = im.getpixel((w, h))
            x = pixel[0]
            y = pixel[1]
            z = pixel[2]
            if count == len(string):
                break
            x = x - mod(x, 2) + int(string[count])
            im.putpixel((w, h), (x, y, z))
            count = count + 1
            if count == len(string):
                break
            y = y - mod(y, 2) + int(string[count])
            im.putpixel((w, h), (x, y, z))
            count = count + 1
            if count == len(string):
                break
            z = z - mod(z, 2) + int(string[count])
            im.putpixel((w, h), (x, y, z))
            count = count + 1
    im.save(new_pic)


def get_key(key_or_path):
    if "BEGIN PUBLIC KEY" in key_or_path or "BEGIN PRIVATE KEY" in key_or_path:
        pem_data = key_or_path
    else:
        with open(key_or_path) as f:
            pem_data = f.read()
    return RSA.importKey(pem_data)


def rsa_encrypt(msg, pub_path, max_length=100):
    """
    RSA加密
    :param msg: 加密字符串
    :param pub_path: 公钥路径
    :param max_length: 1024bit的秘钥不能超过117， 2048bit的秘钥不能超过245
    :return:
    """
    key = get_key(pub_path)
    cipher = PKCS1_cipher.new(key)
    res_byte = bytes()
    for i in range(0, len(msg), max_length):
        res_byte += cipher.encrypt(msg[i:i + max_length].encode(default_encoding))
    return base64.b64encode(res_byte).decode(default_encoding)


def rsa_decrypt(msg, pri_path, max_length=256):
    """
    RSA解密
    :param msg: 加密字符串
    :param pri_path: 私钥路径
    :param max_length: 1024bit的秘钥用128，2048bit的秘钥用256位
    :return:
    """
    key = get_key(pri_path)
    cipher = PKCS1_cipher.new(key)

    res_bytes = bytes()
    encrypt_data = base64.b64decode(msg)
    for i in range(0, len(encrypt_data), max_length):
        res_bytes += cipher.decrypt(encrypt_data[i:i + max_length], 0)
    return res_bytes.decode(default_encoding)


def extract(pic, lenth, hide):
    binary = ""
    string = ""
    count = 0
    im = Image.open(pic)
    width = im.size[0]
    height = im.size[1]
    for h in range(height):
        for w in range(width):
            pixel = im.getpixel((w, h))
            x = pixel[0]
            y = pixel[1]
            z = pixel[2]
            if count == lenth:
                break
            binary = binary + str(mod(x, 2))
            count = count + 1
            if count == lenth:
                break
            binary = binary + str(mod(y, 2))
            count = count + 1
            if count == lenth:
                break
            binary = binary + str(mod(z, 2))
            count = count + 1
    with open(hide, "w", encoding="UTF-8") as f:
        for i in range(0, len(binary), 8):
            string = string + chr(int(binary[i:i + 8], 2))
        f.write(string)


if __name__ == '__main__':
    pic = r"/home/feifan/code/Python/ImageLSB_RSA/image/1.png"  # 获得原始图像位置
    flag = r"/home/feifan/code/Python/ImageLSB_RSA/text/flag.txt"   # 获得需要加密信息
    rsa_encrypt_flag = "/home/feifan/code/Python/ImageLSB_RSA/text/rsa_encrypt_flag.txt"    # rsa加密后信息
    rsa_decrypt_flag = "/home/feifan/code/Python/ImageLSB_RSA/text/rsa_decrypt_flag.txt"  # rsa解密后信息
    new_pic = r"/home/feifan/code/Python/ImageLSB_RSA/image/new.png"    #储存加密图像位置
    pub_position = "/home/feifan/code/Python/ImageLSB_RSA/key/pub_key.pem"  # 公钥位置
    pri_position = "/home/feifan/code/Python/ImageLSB_RSA/key/pri_key.pem"  # 密钥位置
    hidden_msgs = "/home/feifan/code/Python/ImageLSB_RSA/text/hide.txt"
    LSB_msgs = "/home/feifan/code/Python/ImageLSB_RSA/text/LSB.txt"
    key = RSA.generate(2048)   # 生成2048位公钥和密钥
    pri_key = key.export_key()  # 存储密钥
    with open(pri_position, "wb") as f:
        f.write(pri_key)

    pub_key = key.public_key().export_key()  # 存储公钥
    with open(pub_position, "wb") as f:
        f.write(pub_key)

    original_msg = bin_ord(flag)
    encrypted_data = rsa_encrypt(original_msg, pub_position)    # 对原始信息进行rsa加密
    lenth = len(encrypted_data)*8
    print("encrypt_data:", encrypted_data)  # 打出rsa加密后的内容
    with open(rsa_encrypt_flag,'wb') as file:   # 写入加密后的文件
        file.write(encrypted_data.encode())

    hide(pic, rsa_encrypt_flag, new_pic)    # 使用LSB（Least Significant Bit)，选择最低有效位对图像进行隐写
    extract(new_pic, lenth, hidden_msgs)
    with open(hidden_msgs,'r') as file:
        content = file.read()

    print("encrypt_data:", content)
    decrypted_data = rsa_decrypt(content.encode(), pri_position)
    ascii_text = ''.join([chr(int(decrypted_data[i:i + 8], 2)) for i in range(0, len(decrypted_data), 8)])
    with open(rsa_decrypt_flag,'wb') as file:
        file.write(ascii_text.encode())

    print("decrypt_data:", ascii_text)
        # 采用LSB方法对图片进行加密